// APIHook.cpp : Defines the exported functions for the DLL application.
//

#include "stdafx.h"
#include "APIHook.h"
#include "APITypeDef.h"
#include "DllFunInfo.h"
#include "../Detours/include/detours.h"
#include "ShareMemory.h"
#include <vector>


LPShareDllInfo g_pShareInfo=NULL;





using namespace std;



#pragma comment(lib, "detoured.lib")
#pragma comment(lib, "detours.lib")





DWORD g_MyPID=0;
HMODULE g_hDLL;	//HMODULE of this DLL

vector<FunAddrData >g_HookFunData;

BOOL APIENTRY DllMain(HMODULE hDLL, DWORD Reason, LPVOID Reserved)
{
	switch (Reason)
	{
	case DLL_PROCESS_ATTACH:
		{
			g_hDLL=hDLL;
			DisableThreadLibraryCalls(hDLL);
			g_pShareInfo=(LPShareDllInfo)CShareMemory::GetShareMemory();
			CDllFunInfo::InitFunAddr();
			g_HookFunData.clear();
			g_MyPID=GetCurrentProcessId();
			InterlockedIncrement(&g_pShareInfo->m_DllLoadNum);
			CAPIHook::HookAllFun();
			break;
		}
	case DLL_THREAD_ATTACH:
	case DLL_THREAD_DETACH:
	case DLL_PROCESS_DETACH:
		{
			InterlockedDecrement(&g_pShareInfo->m_DllLoadNum);
			CAPIHook::UnHookFun(FunNum);
			CShareMemory::ReleaseMemory();
			break;
		}
		
	}
	return TRUE;
}



// CAPIHook::CAPIHook()
// {
// 	return;
// }

int CAPIHook::GetInjectedDllNum(void)
{
	return g_pShareInfo->m_ProcessHookedNum;
}
int CAPIHook::GetLoadedDllNum(void)
{
	return g_pShareInfo->m_DllLoadNum;
}

DWORD CAPIHook::SetThisProcessFireWall(HWND hwnd)
{
	BOOL isNeedRefresh=FALSE;
	if (g_pShareInfo->m_FireWallExePID!=0)
	{
		if (OpenProcess(NULL,FALSE,g_pShareInfo->m_FireWallExePID)==NULL)
		{
			isNeedRefresh=TRUE;
		}
	}
	if(isNeedRefresh==TRUE || g_pShareInfo->m_FireWallExePID==0)
	{
		g_pShareInfo->m_FireWallExePID=g_MyPID;
		g_pShareInfo->m_FireWallHwnd=hwnd;
	}
	return g_pShareInfo->m_FireWallExePID;
}

BOOL CAPIHook::HookFun( FunNameID funName,BOOL isForceHook )
{
	BOOL retval=FALSE;

	if (g_MyPID!=g_pShareInfo->m_FireWallExePID || isForceHook==TRUE)
	{
		FunAddrData fd={0};
		if (CheckIsHooked(funName)==FALSE)//unHooked
		{
			fd.FunName=funName;
			fd.HookFunAddr=CDllFunInfo::GetHookFunAddress(funName);
			fd.OldFunAddr=DetourFindFunction(CDllFunInfo::GetFunDllName(funName),CDllFunInfo::GetFunStrName(funName));
			if (fd.OldFunAddr!=NULL)
			{
				retval=BeginRedirect(fd);
				if (retval==TRUE)
				{
					InterlockedIncrement(&g_pShareInfo->m_ProcessHookedNum);
				}
			}
		}
		else//Hooked,need rehook
		{
			BeginRedirect(*GetFunAddrData(funName));
		}
		
	}
	return retval;
}

BOOL CAPIHook::BeginRedirect(FunAddrData fd)
{
	BOOL retval=FALSE;


	DetourTransactionBegin();
	DetourUpdateThread(GetCurrentThread());
	DetourAttach(&fd.OldFunAddr,fd.HookFunAddr);
	if (DetourTransactionCommit()==NO_ERROR)
	{
		retval=TRUE;
		if (CheckIsHooked(fd.FunName)==FALSE)
		{
			g_HookFunData.push_back(fd);
		}
	}
	return retval;
}

BOOL CAPIHook::UnHookFun( FunNameID funName,BOOL isForceUnHook)
{
	BOOL retval=FALSE;
	if (g_MyPID!=g_pShareInfo->m_FireWallExePID || isForceUnHook==TRUE)
	{
		if (funName<FunNum)
		{
			LPFUNADDRDATA lpFun=GetFunAddrData(funName);
			if (lpFun!=NULL)
			{
				InterlockedDecrement(&g_pShareInfo->m_ProcessHookedNum);
				
				DetourTransactionBegin();
				DetourUpdateThread(GetCurrentThread());
				DetourDetach(&lpFun->OldFunAddr,lpFun->HookFunAddr);
				DetourTransactionCommit();
				DWORD nID=GetFunAddrID(funName);
				if (nID!=-1)
				{
					g_HookFunData.erase(g_HookFunData.begin()+nID);
				}
			}
			retval=TRUE;
		}
		else if (funName==FunNum)
		{
			int nCount=g_HookFunData.size();
			for (int i=0;i<nCount;i++)
			{

				InterlockedDecrement(&g_pShareInfo->m_ProcessHookedNum);

				DetourTransactionBegin();
				DetourUpdateThread(GetCurrentThread());
				DetourDetach(&g_HookFunData[i].OldFunAddr,g_HookFunData[i].HookFunAddr);
				DetourTransactionCommit();
			}
			g_HookFunData.clear();
			retval=TRUE;
		}
	}
	return retval;
}

BOOL CAPIHook::HookAllFun(BOOL isForceHook)
{
	BOOL retval=FALSE;
	int SuccessCount=0;
	if (isForceHook==TRUE || g_MyPID!=g_pShareInfo->m_FireWallExePID)
	{
		for (int i=0;i<FunNum;i++)
		{
			if (HookFun((FunNameID(i)))==TRUE)
			{
				SuccessCount++;
			}
		}
	}
	if (SuccessCount==FunNum)
	{
		retval=TRUE;
	}
	return retval;
}

BOOL CAPIHook::CheckIsHooked( FunNameID funName )
{
	BOOL retval=FALSE;
	int nCount=g_HookFunData.size();
	for (int i=0;i<nCount;i++)
	{
		if (g_HookFunData[i].FunName==funName)
		{
			retval=TRUE;
			break;
		}
	}
	return retval;
}

BOOL CAPIHook::IsNetBlocked()
{
	return g_pShareInfo->m_BlockNetwork;
}

void CAPIHook::BlockNet( BOOL isBlock)
{
	g_pShareInfo->m_BlockNetwork=isBlock;
}